# angle_v_search_new.py
import pandas as pd
import numpy as np
import importlib.util
import sys
import io
import contextlib
import random
from multiprocessing import Pool, cpu_count
import time


# 加载three_smoke_modified模块
def load_module(file_path: str, name: str = "target_mod"):
    spec = importlib.util.spec_from_file_location(name, file_path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    assert spec.loader is not None
    spec.loader.exec_module(mod)
    return mod


# 加载模块
tsm_mod = load_module("three_smoke_modified.py")


class SuppressToExcel:
    """临时禁用 pandas.DataFrame.to_excel，避免被调用程序每次写盘"""

    def __enter__(self):
        self._orig = pd.DataFrame.to_excel
        pd.DataFrame.to_excel = lambda *args, **kwargs: None
        return self

    def __exit__(self, exc_type, exc, tb):
        pd.DataFrame.to_excel = self._orig


def sweep_v_for_case(drone, missile, alpha_deg, v_start=70.0, v_end=140.0, step=1.0, seed=42):
    """对单个情况(无人机,导弹,角度)进行速度扫描"""
    results = []
    v_vals = np.round(np.arange(v_start, v_end + 1e-9, step), 2)

    for v in v_vals:
        # 固定随机种子，保证不同v的可重复性
        random.seed(seed)
        np.random.seed(seed)

        # 调用优化：静默stdout + 禁写Excel
        with SuppressToExcel(), contextlib.redirect_stdout(io.StringIO()):
            try:
                total, drops, bursts, indiv_times = tsm_mod.main(
                    drone, missile, alpha_deg, v, save_excel=False
                )

                results.append({
                    "drone": drone,
                    "missile": missile,
                    "alpha_deg": alpha_deg,
                    "v": float(v),
                    "total_shield": float(total),
                    "drop_times": drops,
                    "burst_times": bursts,
                    "indiv_times": indiv_times
                })
            except Exception as e:
                print(f"Error for {drone}-{missile} at v={v}: {e}")
                results.append({
                    "drone": drone,
                    "missile": missile,
                    "alpha_deg": alpha_deg,
                    "v": float(v),
                    "total_shield": 0.0,
                    "drop_times": [0, 0, 0],
                    "burst_times": [0, 0, 0],
                    "indiv_times": [0, 0, 0]
                })

        print(f"{drone}-{missile} at alpha={alpha_deg} v={v:.2f} completed, total={results[-1]['total_shield']:.3f}")

    return results


def find_best_v_for_case(case):
    """为单个情况找到最优速度"""
    drone, missile, alpha_deg = case

    print(f"Finding best v for {drone}-{missile} at alpha={alpha_deg}")

    # 粗扫：70-140 m/s，步长1 m/s
    coarse_results = sweep_v_for_case(drone, missile, alpha_deg, 70.0, 140.0, 1.0)
    coarse_df = pd.DataFrame(coarse_results)

    # 找到粗扫最优速度
    best_v_coarse = coarse_df.loc[coarse_df['total_shield'].idxmax(), 'v']
    best_total_coarse = coarse_df.loc[coarse_df['total_shield'].idxmax(), 'total_shield']

    print(f"Coarse scan: best v={best_v_coarse:.2f}, total={best_total_coarse:.3f}")

    # 细扫：最优速度附近±5 m/s，步长0.5 m/s（从0.1增加到0.5）
    fine_start = max(70.0, best_v_coarse - 5.0)
    fine_end = min(140.0, best_v_coarse + 5.0)
    fine_results = sweep_v_for_case(drone, missile, alpha_deg, fine_start, fine_end, 0.5)  # 步长改为0.5
    fine_df = pd.DataFrame(fine_results)

    # 找到细扫最优速度
    best_v_fine = fine_df.loc[fine_df['total_shield'].idxmax(), 'v']
    best_total_fine = fine_df.loc[fine_df['total_shield'].idxmax(), 'total_shield']

    print(f"Fine scan: best v={best_v_fine:.2f}, total={best_total_fine:.3f}")

    # 精扫：最优速度附近±1 m/s，步长0.1 m/s（从0.02增加到0.1）
    precise_start = max(70.0, best_v_fine - 1.0)
    precise_end = min(140.0, best_v_fine + 1.0)
    precise_results = sweep_v_for_case(drone, missile, alpha_deg, precise_start, precise_end, 0.1)  # 步长改为0.1
    precise_df = pd.DataFrame(precise_results)

    # 找到精扫最优速度
    best_v_precise = precise_df.loc[precise_df['total_shield'].idxmax(), 'v']
    best_total_precise = precise_df.loc[precise_df['total_shield'].idxmax(), 'total_shield']

    print(f"Precise scan: best v={best_v_precise:.2f}, total={best_total_precise:.3f}")

    # 返回最佳结果
    best_result = precise_df.loc[precise_df['total_shield'].idxmax()].to_dict()

    return best_result


if __name__ == "__main__":
    # 读取已有的角度优化结果
    angle_df = pd.read_excel("angle_total.xlsx")

    # 筛选出有效情况（total_shield > 0）
    valid_cases = angle_df[angle_df['total_shield'] > 0]

    # 准备参数列表
    cases = [(row['drone'], row['missile'], row['alpha_deg']) for _, row in valid_cases.iterrows()]

    print(f"Found {len(cases)} valid cases to optimize for velocity")

    # 使用多进程并行处理
    start_time = time.time()

    with Pool(processes=cpu_count()) as pool:
        best_results = pool.map(find_best_v_for_case, cases)

    # 创建结果DataFrame
    best_v_df = pd.DataFrame(best_results)

    # 合并原有角度优化结果和速度优化结果
    merged_df = angle_df.merge(
        best_v_df[['drone', 'missile', 'v', 'total_shield', 'drop_times', 'burst_times', 'indiv_times']],
        on=['drone', 'missile'],
        how='left',
        suffixes=('_angle', '_v')
    )

    # 更新总遮蔽时间为速度优化后的值
    merged_df['total_shield'] = merged_df['total_shield_v'].fillna(merged_df['total_shield_angle'])

    # 更新投放时间和引信延时
    for i in range(3):
        merged_df[f'drop_time_{i + 1}'] = merged_df['drop_times_v'].apply(
            lambda x: x[i] if isinstance(x, list) and len(x) > i else np.nan)
        merged_df[f'burst_time_{i + 1}'] = merged_df['burst_times_v'].apply(
            lambda x: x[i] if isinstance(x, list) and len(x) > i else np.nan)
        merged_df[f'indiv_time_{i + 1}'] = merged_df['indiv_times_v'].apply(
            lambda x: x[i] if isinstance(x, list) and len(x) > i else np.nan)

    # 删除不必要的列
    merged_df = merged_df.drop(['drop_times_angle', 'burst_times_angle', 'indiv_times_angle',
                                'drop_times_v', 'burst_times_v', 'indiv_times_v',
                                'total_shield_angle', 'total_shield_v'], axis=1, errors='ignore')

    # 保存结果
    merged_df.to_excel("angle_v_total.xlsx", index=False)

    end_time = time.time()
    print(f"所有计算完成，总耗时: {end_time - start_time:.2f}秒")
    print("结果已保存到 angle_v_total_new.xlsx")